import os
import shutil as st

def divideData(path, traindst, testdst):
    packages = [path + "/" + i + "/" for i in os.listdir(path)]
    trainDstDirs = [traindst + "/" + i + "/" for i in os.listdir(path)]
    testDstDirs = [testdst + "/" + i + "/" for i in os.listdir(path)]

    for package, trainDst, testDst in zip(packages, trainDstDirs, testDstDirs):
        imgs = [ i for i in os.listdir(package)]
        length, cnt = len(imgs), 0
        train = int(length * 0.7)
        test = length - train

        for img in imgs:
            if cnt < train:
                if not os.path.exists(trainDst):
                   os.makedirs(trainDst)
                st.copy(package + img, trainDst + img)
                print(">> " + trainDst + img + " is OK")
                cnt += 1

            else:
                if not os.path.exists(testDst):
                   os.makedirs(testDst)
                st.copy(package + img, testDst + img)
                print(">> " + testDst + img + " is OK")
                cnt += 1

if __name__ == "__main__":
    divideData("./picture", "./train", "./test")
